package top.jfunc.common.db.data;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.util.ArrayList;
import java.util.List;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.IntConsumer;

import javax.sql.DataSource;

import top.jfunc.common.utils.IoUtil;

/**
 * @author <a href="https://github.com/TAKETODAY">Harry Yang</a>
 * @since 1.0 2022/11/12 11:02
 */
final class SimpleDataExtractor extends DataExtractor {
  private static final Logger logger = LoggerFactory.getLogger(SimpleDataExtractor.class);

  private final String querySQL;

  // @Nullable
  private final List<Function<String, String>> sqlModifiers;

  SimpleDataExtractor(String querySQL, DataSource dataSource,
      Consumer<Throwable> exceptionConsumer,
      PreparedStatementSetter statementPostProcessor,
      List<Function<String, String>> sqlModifiers, IntConsumer successCallback) {
    super(dataSource, exceptionConsumer, statementPostProcessor, successCallback);
    this.querySQL = querySQL;
    this.sqlModifiers = sqlModifiers;
  }

  public String getQuerySQL() {
    return querySQL;
  }

  @Override
  public int processRow(RowCallbackHandler handler) {
    int rowNumber = 0;
    ResultSet resultSet = null;
    Connection connection = null;
    PreparedStatement statement = null;

    try {
      String querySQL = getQuerySQL();
      if (sqlModifiers != null) {
        for (Function<String, String> sqlModifier : sqlModifiers) {
          querySQL = sqlModifier.apply(querySQL);
        }
      }

      connection = dataSource.getConnection();
      statement = prepareStatement(connection, querySQL);
      configStatement(statement);
      resultSet = statement.executeQuery();

      while (resultSet.next()) {
        handler.processRow(resultSet, ++rowNumber);
      }

      handler.finishedReading(rowNumber);

      if (successCallback != null) {
        successCallback.accept(rowNumber);
      }
      return rowNumber;
    }
    catch (Throwable e) {
      if (exceptionConsumer != null) {
        exceptionConsumer.accept(e);
        return rowNumber;
      }
      else {
        throw new IllegalStateException("数据读取失败", e);
      }
    }
    finally {
      IoUtil.closeAll(statement,resultSet,connection);
    }
  }

  public static Builder builder() {
    return new Builder();
  }

  public static Builder fetch(String querySQL, DataSource dataSource) {
    return new Builder()
        .querySQL(querySQL)
        .dataSource(dataSource);
  }

  //

  static final class Builder extends DataExtractorBuilderSupport<Builder>
      implements DataExtractor.SimpleBuilder {

    private String querySQL;
    private final List<Function<String, String>> sqlModifiers = new ArrayList<>();

    @Override
    protected Builder self() {
      return this;
    }

    @Override
    public Builder querySQL(String querySQL) {
      this.querySQL = querySQL;
      return this;
    }

    /**
     * 可以设置多个修改器
     */
    @Override
    public Builder modifySQL(Function<String, String> sqlModifier) {
      this.sqlModifiers.add(sqlModifier);
      return this;
    }

    @Override
    public SimpleDataExtractor build() {
      return new SimpleDataExtractor(querySQL, dataSource,
          exceptionConsumer, statementPostProcessor, sqlModifiers, successCallback);
    }

  }

}
